
import struct
import socket
from packet import Packet

IP_ADDR_SIZE = 4
IP_VERSION_AND_SIZE = 0x45
IP_HDR_SIZE = 20

class IP(Packet):
    def __init__(self, src_ip, dst_ip, pkt_id, protocol, data, ttl=128):
        Packet.__init__(self)
        
        # check packet arguments
        if (IP_ADDR_SIZE != len(src_ip)) or (IP_ADDR_SIZE != len(dst_ip)):
            raise Exception("Invalid IP address length")
        
        # set local members
        self.m_src_ip = src_ip
        self.m_dst_ip = dst_ip
        self.m_id = pkt_id
        self.m_data = data
        
        # build headers
        self.m_header = struct.pack("!BBH", IP_VERSION_AND_SIZE, 0, len(self.m_data) + IP_HDR_SIZE)
        self.m_header += struct.pack("!HH", self.m_id, 0)           # (ID, fragmentation offset)
        self.m_header += struct.pack("!BBH", ttl, protocol, 0)       # (TTL, Protocol, Checksum)
        self.m_header += self.m_src_ip + self.m_dst_ip
        
        self._fix_hdr_checksum()
    
    @staticmethod
    def _chksum_break_to_16bit(hdr_data):
        temp = map(ord, hdr_data)
        return sum(map(lambda x: (x[0] << 8) + x[1], zip(temp[::2], temp[1::2])))
        
    def _fix_hdr_checksum(self):
        chksum = IP._chksum_break_to_16bit(self.m_header)
        while chksum > 0xFFFF:
            chksum_hex = hex(chksum)[2:]
            chksum_hex = (-(len(chksum_hex) % -4) * "0") + chksum_hex
            chksum = IP._chksum_break_to_16bit(chksum_hex.decode("hex"))
            
        chksum = ~long(chksum)
        self.m_header = self.m_header[:10] + struct.pack("!H", chksum & 0xFFFF) + self.m_header[12:]
    
    @staticmethod
    def ip_addr_aton(ip_addr_str):
        return socket.inet_aton(ip_addr_str)